package com.gaga.console.base.utils;

/**
 * @DESCRIPTION:
 * @USER: ningxiang
 * @DATE: 5/3/2025 12:00 AM
 */

import com.alibaba.fastjson.JSON;
import com.gaga.console.core.entity.vo.GPUInfoVO;
import jcuda.driver.CUdevice;
import jcuda.driver.JCudaDriver;
import lombok.extern.slf4j.Slf4j;

import static jcuda.driver.JCudaDriver.*;

@Slf4j
public class GPUInfoUtils {

    // 获取显存大小（单位：GB）和驱动版本
    public static GPUInfoVO getGPUInfo() {

        GPUInfoVO gpuInfoVO = new GPUInfoVO();

        // 初始化CUDA驱动
        JCudaDriver.setExceptionsEnabled(true);
        try {
            cuInit(0);
        } catch (Exception e) {
            log.error("CUDA初始化失败: " + e.getMessage());
        }

        // 获取CUDA驱动版本
        int[] driverVersion = new int[1];
        cuDriverGetVersion(driverVersion);
        log.error("CUDA驱动版本: %d.%d\n",
                driverVersion[0] / 1000, (driverVersion[0] % 1000) / 10);

        gpuInfoVO.setGpuDriverVersion(String.valueOf(driverVersion[0] / 1000));

        // 获取设备数量
        int[] deviceCount = new int[1];
        cuDeviceGetCount(deviceCount);
        log.error("发现CUDA设备数量: " + deviceCount[0]);

        // 遍历设备并获取显存
        for (int i = 0; i < deviceCount[0]; i++) {
            CUdevice device = new CUdevice();
            cuDeviceGet(device, i);

            // 获取设备名称
            byte[] deviceName = new byte[256];
            cuDeviceGetName(deviceName, deviceName.length, device);
            String name = new String(deviceName).trim();

            // 获取显存大小（字节）
            long[] memSize = new long[1];
            cuDeviceTotalMem(memSize, device);
            long gb = memSize[0] / (1024 * 1024 * 1024);
            log.error("显存大小: " + gb + " GB");
            gpuInfoVO.setGpuMemorySize(String.valueOf(gb));

            log.error("设备 %d: %s, 显存: %d GB\n", i, name, gb);
        }
        return gpuInfoVO;
    }

    public static void main(String[] args) {
        System.out.println(JSON.toJSONString(getGPUInfo()));
    }
}
